package simpletrace

import (
	"net/http"
)

func SimpleTraceHTTPMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ctx := r.Context()
		var traceID string
		traceID = r.Header.Get(HeaderSimpleTraceID)
		if traceID == "" {
			next.ServeHTTP(w, r)
			return
		}

		isEnable := r.Header.Get(HeaderSimpleTraceEnable)
		if isEnable == "" {
			next.ServeHTTP(w, r)
			return
		}

		spanID, partentID, region := GetSpanIDAndParentIDFromRequest(r)

		if spanID == "" && partentID == "" {
			logger.WarnCtx(ctx, "[SimpleTraceMiddleware] Simple trace already enabled for %s, traceID: %s, parentID: %s, spanID: %s, region: %s", r.URL.RawPath, traceID, partentID, spanID, region)
			next.ServeHTTP(w, r)
			return
		}

		span, ctx := StartSpan(ctx, r.URL.RawPath, map[string]any{
			"method": r.Method,
			"url":    r.URL.String(),
		}, true, WithIDs(traceID, spanID, partentID))

		r = r.WithContext(ctx)

		next.ServeHTTP(w, r)

		span.Finish(ctx, map[string]any{}, false)
	})
}

func GetSpanIDAndParentIDFromRequest(req *http.Request) (string, string, string) {
	spanID := req.Header.Get(HeaderSimpleTraceSpanID)
	if spanID == "" {
		spanIDCookie, _ := req.Cookie(HeaderSimpleTraceSpanID)
		if spanIDCookie != nil {
			spanID = spanIDCookie.Value
		}
	}
	parentID := req.Header.Get(HeaderSimpleTraceParentID)
	if parentID == "" {
		parentIDCookie, _ := req.Cookie(HeaderSimpleTraceParentID)
		if parentIDCookie != nil {
			parentID = parentIDCookie.Value
		}
	}
	region := req.Header.Get(HeaderSimpleTraceRegion)
	if region == "" {
		regionCookie, _ := req.Cookie(HeaderSimpleTraceRegion)
		if regionCookie != nil {
			region = regionCookie.Value
		}
	}
	return spanID, parentID, region
}
